from PyQt5.QtCore import QThread, pyqtSignal
from yt_dlp import YoutubeDL
from webvtt import WebVTT
import os
import re

class SrtThread(QThread):
    begin_download = pyqtSignal(object)
    download_result = pyqtSignal(object, bool, str)

    def __init__(self, parent, task):
        super().__init__(parent)
        self.task = task

    def run(self) -> None:
        output_dir = self.task.path
        os.makedirs(output_dir, exist_ok=True)

        base_filename = self.sanitize_filename(self.task.title)
        vtt_path = os.path.join(output_dir, base_filename + '.zh-Hans.vtt')
        srt_path = os.path.join(output_dir, base_filename + '.srt')

        ydl_opts_base = {
            'skip_download': True,
            'subtitlesformat': 'vtt',
            'outtmpl': os.path.join(output_dir, base_filename + '.%(ext)s'),
            'quiet': False,
            'http_headers': {
                'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 '
                              '(KHTML, like Gecko) Chrome/123.0.0.0 Safari/537.36',
            },
            'cookiefile': 'cookies.txt',
            'retries': 3
        }

        # 检查是否有字幕
        with YoutubeDL(ydl_opts_base) as ydl:
            info = ydl.extract_info(self.task.url, download=False)
            subtitles = info.get('subtitles', {})
            auto_subtitles = info.get('automatic_captions', {})

            if 'zh-Hans' in subtitles:
                ydl_opts_base.update({
                    'writesubtitles': True,
                    'subtitleslangs': ['zh-Hans'],
                })
            elif 'zh-Hans' in auto_subtitles:
                ydl_opts_base.update({
                    'writeautomaticsub': True,
                    'subtitleslangs': ['zh-Hans'],
                })
            else:
                print(f"❌ 无字幕：{self.task.url}")
                self.download_result.emit(self.task, False, '当前视频没有字幕')
                return

        # 下载字幕
        with YoutubeDL(ydl_opts_base) as ydl:
            self.begin_download.emit(self.task)
            try:
                ydl.download([self.task.url])
            except:
                pass

        # 转换字幕
        self.convert_vtt_to_srt(vtt_path, srt_path)

    def convert_vtt_to_srt(self, vtt_path, srt_path):
        try:
            WebVTT().read(vtt_path).save_as_srt(srt_path)
            os.remove(vtt_path)
            self.download_result.emit(self.task, True, '下载字幕成功')
        except Exception as e:
            print(f"字幕转换失败: {vtt_path} -> {srt_path}, 原因: {e}")
            self.download_result.emit(self.task, False, '转换失败')

    @staticmethod
    def sanitize_filename(filename):
        # Windows 不允许的字符替换成 _
        return re.sub(r'[\\/:*?"<>|]', '_', filename)
